#include <iostream>
#include <iomanip>
#include <fstream>
#include <stdexcept>
#include <queue>
#include <utility>
#include <vector>
#include <sstream>
#include <string>
#include <map>
#include <set>
#include "pasex.hh"

#define ENABLE_DEBUG 1

#if ENABLE_DEBUG
#	define DEBUG(X...) X
#else
#	define DEBUG(X...)
#endif

namespace lex
{
	std::istream & operator>>(std::istream & is, symbol & s)
	{
		std::string str;
		is >> str;
		symbol tmp(str);
		s = std::move(tmp);
		return is;
	}

	bool operator<(symbol const& a, symbol const& b) { return a.s < b.s; }


	std::istream & operator>>(std::istream & is, regular_rule & rp)
	{
		is >> rp.left;
		for (symbol tmp; is >> tmp; ) rp.right.push_back(tmp);
		return is;
	}


	void regular_grammar::push(regular_rule const & rp)
	{
		rps.insert({rp.left, rp});
		if (rp.left.mark && not_mark(rp.left)) marks.push_back(rp.left);
		if (rp.right[0].terminal)
			for (auto ch : rp.right[0].s) alpha.insert(ch);
	}

	bool regular_grammar::not_mark(symbol const & s) const
	{
		for (auto i : marks) if (i.s == s.s) return false;
		return true;
	}

	std::set<symbol> regular_grammar::find_all_termianl(symbol const & syb) const
	{
		std::set<symbol> ret;
		if (syb.terminal) { ret.insert(syb); return ret; }
		auto range = rps.equal_range(syb);
		for (auto rule = range.first ; rule != range.second; ++rule) {
			auto rr = rule->second;
			auto ts = find_all_termianl(rr.right[0]);
			for (auto tsyb : ts) ret.insert(tsyb);
		}
		return ret;
	}

	void regular_grammar::print()
	{
		std::cout << "regular grammar =\n";
		std::cout << "alpha set: ";
		for (auto ch : alpha) std::cout << ch << ' ';
		std::cout << '\n';

		std::cout << "marks:\n";
		for (auto s : marks) std::cout << s.s << '\n';
		std::cout << '\n';

		std::cout << "regular rules:\n";
		for (auto r : rps) {
			std::cout << r.first.s << "->";
			for (auto ts : r.second.right) std::cout << ts.s << ' ';
			std::cout << '\n';
		}
	}

	std::istream & operator>>(std::istream & is, regular_grammar & g)
	{
		std::string line;
		while (std::getline(is, line)) {
			if (line.size() == 0) continue;
			std::stringstream buf(line);
			regular_rule rp;
			buf >> rp;
			g.push(rp);
		}
		return is;
	}

	namespace nfa
	{
		nfa::nfa(regular_grammar const & g)
		{
			alloc = 0; start = alloc_node(false);
			marks = g.marks;
			alpha = g.alpha;
			vis.clear();
			for (auto ma : g.marks) node_index(ma);
			for (auto ma : g.marks) {
				auto id = node_index(ma);
				find_and_add_edge(ma, g, id, start);
			}
		}

		std::pair<nfa::state, int> nfa::closure(state const & s) const // return type state, flag
		{
			std::pair<nfa::state, int> ret;
			std::set<int> visited;
			std::queue<int> q;
			for (auto i : s) {
				q.push(i);
				visited.insert(i);
				ret.first.insert(i);
				if (nodes[i].is_terminal) ret.second = nodes[i].flag;
			}
			while (!q.empty()) {
				auto now = q.front(); q.pop();
				for (auto i : graph[now]) {
					auto c = i.first;
					auto v = i.second;
					if (c) continue;
					if (visited.find(v) != visited.end()) continue;
					q.push(v);
					visited.insert(v);
					ret.first.insert(v);
					if (nodes[v].is_terminal) ret.second = true;
				}
			}
			return ret;
		}

		nfa::state nfa::trans(nfa::state const & s, char ch) const
		{
			nfa::state ret;
			for (auto i : s) {
				auto range = graph[i].equal_range(ch);
				for (auto it = range.first; it != range.second; ++it) {
					auto v = it->second;
					ret.insert(v);
				}
			}
			return ret;
		}

		void nfa::print(std::ofstream & of)
		{
			of << "digraph\n{\n\tgraph[rankdir=LR]\n";
			for (int i = 0; i < (int)graph.size(); i++)
				for (auto j : graph[i]) {
					std::string ch{j.first};
					if (!j.first) ch = "epsilon";
					of << "\t\"" << i;
					if (nodes[i].is_terminal) of << "* "  << marks[nodes[i].flag - 1].s;
					of << "\" -> \"" << j.second;
					if (nodes[j.second].is_terminal) of << "* " << marks[nodes[j.second].flag - 1].s;
					of << "\" [label=\"" << ch << "\"]\n";
				}
			of << "}\n";
		}

	//private:

		int nfa::alloc_node(bool terminal, int flag)
		{
			node tmp(alloc++, terminal, flag);
			nodes.push_back(tmp);
			std::multimap<char, int> ms;
			graph.push_back(ms);
			return alloc - 1;
		}

		int nfa::node_index(symbol const & s)
		{
			if (symbol_index.find(s) == symbol_index.end())
				symbol_index[s] = alloc_node(false);
			return symbol_index[s];
		}

		void nfa::find_and_add_edge(symbol const & ma, regular_grammar const & g, int flag, int now, symbol const & s)
		{
			int tnow = node_index(ma);
			vis.insert(tnow);
			add_edge(now, tnow, s.s);
			auto range = g.rps.equal_range(ma);
			for (auto rule = range.first; rule != range.second; ++rule) {
				auto rr = rule->second;
				auto tset = g.find_all_termianl(rr.right[0]);
				for (auto const & tsyb : tset) {
					auto ts = tsyb.s;
					if (rr.right.size() == 1) {
						auto id = node_index(rr.left);
						if (rr.right[0].s == "&") {
							nodes[id].is_terminal = true;
							nodes[id].flag = flag;
							continue;
						}
						auto id2 = alloc_node(true, flag);
						add_edge(id, id2, tsyb.s);
					} else {
						auto id = node_index(rr.left);
						auto id2 = node_index(rr.right[1]);
						if (vis.find(id2) == vis.end())
							find_and_add_edge(rr.right[1], g, flag, id, tsyb);
						else
							add_edge(id, id2, tsyb.s);
					}
				}
			}
		}

		void nfa::add_edge(int u, int v, std::string const & s) // "" for epsilon, and has length large than 1's terminal
		{
			if (s == "") { graph[u].insert({0, v}); return; }
			if (s.size() == 1) { graph[u].insert({s[0], v}); return; }
			int now = u, prev = u;
			for (int i = 0; i < (int)s.size() - 1; i++) {
				now = alloc_node(false);
				graph[prev].insert({s[i], now});
				prev = now;
			}
			graph[prev].insert({s[s.size() - 1], v});
		}

		symbol const nfa::epsilon("");
	}


	namespace dfa
	{
		dfa::dfa(nfa::nfa const & n)
		{
			alloc = 0; alpha = n.alpha;
			marks = n.marks;
			auto closure = n.closure({n.start});
			start = node_index(closure.first, closure.second);
			now = start;
			for (int i = 0; i < (int)nodes.size(); i++) {
				for (auto ch : alpha) {
					auto state = n.trans(nodes[i].nfa_id, ch);
					if (state.empty()) continue;
					auto p = n.closure(state);
					auto id = node_index(p.first, p.second);
					add_edge(i, id, ch);
				}
			}
		}

		void dfa::lex_init()
		{
			first = true;
			last_ac = -1; len = 0;
			now = start;
		}

		dfa::token_type dfa::lex(std::string const & str, std::string::iterator & it)
		{
			if (first) { DEBUG(std::cout << "\ntoken table:\n"); first = false; }
			if (it == str.end()) return std::pair<std::string, std::string>("#", "#");
			std::string tok;
			for (; it != str.end(); ++it) {
				if (graph[now].find(*it) == graph[now].end()) {
					if (last_ac == -1) throw std::runtime_error("cannot match: " + tok + "!");
					std::pair<std::string, std::string> ret(
							marks[nodes[last_ac].flag - 1].s,
								std::string(tok.begin(), tok.begin() + len));

					DEBUG(std::cout << "\e[0;35m" << std::left << std::setw(17) << ret.first << "\t\e[1;34m" << ret.second << "\e[0m\n");

					++it;
					if (it == str.end()) return ret;
					it -= tok.size() - len + 1;
					last_ac = -1;
					now = start;
					tok = "";
					return ret;
				} else {
					now = graph[now].at(*it);
					tok += *it;
					if (nodes[now].is_terminal) {
						last_ac = now;
						len = tok.size();
					}
				}
			}
			return std::pair<std::string, std::string>();
		}

		void dfa::lex(std::string const & str)
		{
			dfa::lex_init();
			auto code = str + char(0);
			auto it = code.begin();
			while (std::next(it) != code.end())
				dfa::lex(str, it);
		}

		void dfa::print(std::ofstream & of)
		{
			of << "digraph\n{\n\tgraph[rankdir=LR]\n";
			for (int i = 0; i < (int)graph.size(); i++)
				for (auto j : graph[i]) {
					std::string ch{j.first};
					if (!j.first) ch = "epsilon";
					of << "\t\"" << i;
					if (nodes[i].is_terminal) of << "* " << marks[nodes[i].flag - 1].s;
					of << "\" -> \"" << j.second;
					if (nodes[j.second].is_terminal) of << "* " << marks[nodes[j.second].flag - 1].s;
					of << "\" [label=\"" << ch << "\"]\n";
				}
			of << "}\n";
		}


	// private:
		int dfa::alloc_node(state const & nfa_s, int flag)
		{
			node tmp(nfa_s, alloc++, flag > 0, flag);
			nodes.push_back(tmp);
			std::map<char, int> m;
			graph.push_back(m);
			return alloc - 1;
		}

		int dfa::node_index(state const & s, int flag)
		{
			if (state_index.find(s) == state_index.end())
				state_index[s] = alloc_node(s, flag);
			return state_index[s];
		}
	}
}

#undef DEBUG
#undef ENABLE_DEBUG

/*
int main()
{
//	std::ifstream fin{"test-input"};
	std::ifstream fin{"input-grammar"};
	std::ifstream pin{"input-program"};
	std::ifstream tpin{"test-input-program"};
	std::ofstream nfa_out{"nfa-out.dot"};
	std::ofstream dfa_out{"dfa-out.dot"};
	std::ofstream token_table{"token-table"};
	auto g = regular_grammar::from_istream(fin);
	fin >> g;
	g.print();

	nfa n{g}; // prevent most-vexing parse
	n.print(nfa_out);

	dfa d{n};
	d.print(dfa_out);
//	dfa d{nfa{g}};

	std::string s;
	std::string t;
	for (; pin >> t; s += t) if (s.size()) s +="^";
	d.lex(s, token_table);
}
*/

